# db.py
# by Sean A. Harrison <sah@tyndale.com>
# ---------------------------------------------------------------------------
# $Id: db.py 103 2013-09-03 19:10:18Z sah.harrison@gmail.com $

"""
A simple object-relational mapper that stays out of your way.
Designed to be small, fast, and transparent. Integrates with user-defined record classes.

Classes:
    Record:     base class for user-defined Record classes (inherits from amplitude.dict.Dict)
    RecordSet:  set of Records (inherits from list)
    Database :  a database connection
        (init with a standard DB-API 2.0 connection string)

---------------------------------------------------------------------------
Memory Usage:                                  Footprint:
> python                           #   2240 K     2240 K  (Python 2.6.4 win32)
>>> import psycopg2 as dba         #   2940 K      700 K
>>> from amplitude import db       #   3016 K       76 K

---------------------------------------------------------------------------
Sample session: Try typing what is on the lines beginning with >>>
(Note: requires config.py and config_test.ini with the database connection info.
(Using assignment and comparison in the examples, for doctest happiness)

>>> from amplitude.config import SiteConfig
>>> cf = SiteConfig('config_test.ini')
>>> exec("import %s as dba" % cf.value('Database', 'adaptor'))
>>> d = Database(cf.value('Database', 'connection'), dba)
>>> d.execute("create table table1 (name varchar primary key, email varchar not null unique);")
>>> d.execute("insert into table1 (name, email) values ('sah', 'sah@tyndale.com')")
>>> d.execute("insert into table1 (name, email) values (%(name)s, %(email)s)", {'name':'pat', 'email':'gpl@tyndale.com'})
>>> rs = d.select("select * from table1")
>>> rs == [{'name': 'sah', 'email': 'sah@tyndale.com'}, {'name': 'pat', 'email': 'gpl@tyndale.com'}]
True
>>> rs[0]['name']
'sah'
>>> d.execute("update table1 set email='harrison@tbc.net' where name='sah'")
>>> r = d.select_one("select * from table1 where name='sah' limit 1")
>>> r == {'name': 'sah', 'email': 'harrison@tbc.net'}
True
>>> r['email']
'harrison@tbc.net'
>>> class MyRecordClass(Record): pass           # -- Here's how you can use your own Record classes. --
...
>>> r = d.select_one("select * from table1 where name='sah'", record_class=MyRecordClass)
>>> r.__class__ == MyRecordClass
True
>>> d.execute("drop table table1")
>>> d.connection.close()
"""

import datetime, re, time
from amplitude.dict import Dict
from amplitude.string import identifier, nameify, to_unicode
import sys

class Record(Dict):
    """A record from the database."""

    def __init__(self, db, **args):
        self.__dict__['db'] = db
        for k in args.keys():
            if type(args[k]) == str:
                self[k] = unicode(args[k], 'utf-8')
            else:
                self[k] = args[k]

    # allow dot-notation attribute access alongside dict-notation
    def __getattr__(self, name):
        val = dict.get(self, name)
        if type(val)==str:
            return unicode(val, 'utf-8')
        else:
            return val       # returns None if the key is not found.

    def key_tuple(self, key=None):
        """can be used as a dict key."""
        return tuple([self[k] for k in key or self.pk_fields])
        
    def key_dict(self, key=None):
        """return a copy of this record that only has the given keys. Useful for generating a keyword-argument dictionary for select()"""
        d = self.__class__(self.db)
        key = key or self.pk_fields
        for k in key:
            d[k] = self[k]
        return d

    def as_element(self, fields=[], attribs=[], tag=None):
        """Return the contents of this record as etree Element. If fields != [], use those fields.
        The attribs arguments lists those fields that should be attributes rather than tags; default=pk_fields
        """
        from lxml import etree
        if tag is None: tag = self.__class__.__name__
        tree = etree.Element(tag)
        
        fields = fields[:]      # copy so that manipulations don't change the source lists.
        attribs = attribs[:]
        
        # make attributes of the element
        if attribs == [] and self.pk_fields is not None:
            attribs = self.pk_fields
        
        # make other fields into subelements (and don't include as attribs).
        if fields == []:
            fields = self.keys()
        for a in attribs:
            if a in fields:
                fields.pop(fields.index(a))
                
        # remove null elements from fields and attribs
        for i in range(len(fields)-1, -1, -1):          # count backward to avoid disruptive pop()
            if self[fields[i]] is None:
                fields.pop(i)
        for i in range(len(attribs)-1, -1, -1):
            if self[attribs[i]] is None:
                attribs.pop(i)
        
        # build XML for attribs and fields
        for k in attribs:
            tree.attrib[k] = to_unicode(self[k])
        for k in [key for key in fields if key not in attribs]:
            if self[k] is not None:
                xml = u"<%s>%s</%s>" % (k, to_unicode(self[k]), k)
                e = etree.fromstring(xml, etree.XMLParser(recover=True, remove_blank_text=False))    # make it forgiving.
                tree.append(e)
        
        return tree

    def as_json(self, fields=None):
        """Return the contents of this record as json"""
        import json
        d = Dict()
        for k in fields or self.keys(): 
            d[k] = self[k]
            if type(d[k]) == datetime.datetime:         # datetime not compatible with json.dumps
                d[k] = d[k].isoformat(' ')              # -- iso formatted string
        return json.dumps(d)

    def from_element(self, elem, tag=None):
        """build a record from an ElementTree element"""
        from lxml import etree
        s = self.__class__(self.db)
        if tag==None: tag = nameify(self.__class__.__name__)
        for k in elem.attrib.keys():
            s[k] = elem.get(k)
        for e in elem.getchildren():
            # put the entire contents of the child element into the record as unicode text.
            utxt = etree.tostring(e, encoding='UTF-8').decode('UTF-8').strip()
            utxt = re.sub("^<[^>]*/?>", "", utxt, re.U)         # strip the start and end tag from each child element
            utxt = re.sub("</[^>]*>$", "", utxt, re.U)
            key = re.sub("\{[^\}]*\}", '', e.tag)               # strip the namespace from the tag to form the field name
            s[key] = re.sub(' xmlns(?::[^=]*)="[^"]*"', "", utxt)
        for k in s.keys():
            if s[k] is None: s[k] = ''
            s[k] = s[k].replace("&lt;", '<').replace("&gt;", '>')
        return s        

    def from_prefix(self, prefix, **args):
        """return an instance of this Model containing the keys from self that have the given prefix, with the prefix removed"""
        r = self.__class__.__init__(self, self.db, **args)
        for k in [k for k in self.keys() if k[:len(prefix)]==prefix]:
            r[k[len(prefix):]] = self[k]
        return r

class RecordSet(list):
    """A set of database records"""

    def select(self, **args):
        """return all records in the set that have the attributes specified in args, or empty RecordSet"""
        result = RecordSet()
        for record in self:
            match = True
            for k in args:
                if record[k] != args[k]:
                    match = False
                    break
            if match == True:
                result.append(record)
        return result
        
    def select_one(self, **args):
        """return the first record in the set that has the attributes specified in args, or None."""
        for record in self:
            match = True
            for k in args:
                if record[k] != args[k]:
                    match = False
                    break
            if match == True:
                return record

    def as_dict(self, key):
        """returns the RecordSet as a dict of records, keyed to strings or tuples. 
            keys must be unique, or later duplicates will overwrite earlier ones."""
        d = Dict()
        for record in self:
            if type(key) in [str, unicode]:
                # string key
                d[record[key]] = record
            else:
                # tuple key
                d[record.key_tuple(key)] = record
        return d
        
    def as_element(self, root=None, tag=None, fields=[], attribs=[]):
        from lxml import etree
        if root is None: root = self.__class__.__name__
        doc = etree.Element(root)
        for record in self:
            doc.append(record.as_element(fields=fields, attribs=attribs, tag=tag))
        return doc
        
    def as_xml(self, root=None, encoding='UTF-8', fields=[], attribs=[]):
        from lxml import etree
        doc = self.as_element(root=root, fields=fields, attribs=attribs)
        if encoding=='unicode':
            return etree.tostring(doc, encoding=encoding, pretty_print=True)
        else:
            return etree.tostring(doc, encoding=encoding, xml_declaration=True, pretty_print=True)

    def as_json(self, fields=[]):
        """Return the contents of this record set as json"""
        l = [record.as_json(fields=fields) for record in self]
        return u'[' + u', '.join(l) + u']'

        
class Database(Dict):
    """a database connection object."""

    def __init__(self, __connection_string__=None, dba=None, tries=3, **args):
        Dict.__init__(self, **args)
        self.DEBUG = self.config and self.config.Site and self.config.Site.debug and eval(self.config.Site.debug) or False
        if dba == None:
            import sqlite3 as dba
        elif type(dba) in (str, unicode):
            exec "import %s as dba" % dba
        self.dba = dba
        if dba.__name__ == 'psycopg2':
            # make psycopg2 always return unicode strings
            try:
                dba.extensions.register_type(dba.extensions.UNICODE)
                dba.extensions.register_type(dba.extensions.UNICODEARRAY)                    
            except:
                # if that didn't work for some reason, then just go with the default setup.
                pass
            
        # try reaching the db "tries" times, with increasing wait times, before raising an exception.
        for i in range(tries):
            try: 
                if __connection_string__ != None:
                    self.connection = dba.connect(__connection_string__)
                else:
                    self.connection = dba.connect(**args)
                break   # succeeded, no need to try again
            except: 
                if i==range(tries)[-1]:
                    # last try failed
                    raise
                else:
                    # wait a bit
                    if self.DEBUG == True: 
                        print "  database unreachable, waiting", 2*i, "seconds"
                    time.sleep(2*i)

    def cursor(self):
        """get a cursor for fine-grained transaction control."""
        cursor = self.connection.cursor()
        return cursor

    def execute(self, sql, vals=None, cursor=None):
        """execute SQL transaction, commit it, and return nothing. If a cursor is specified, work within that transaction."""
        if self.DEBUG == True: print '--db.execute(sql):--\n' + sql.strip().encode('utf-8') + '\n'
        try:
            c = cursor or self.connection.cursor()
            if vals in [None, (), [], {}]:
                c.execute(sql)
            else:
                c.execute(sql, vals)
            if cursor is None:
                self.commit()
        except:
            if self.DEBUG==True: print '-- ROLLBACK --'
            self.rollback()
            raise

    def commit(self):
        """commit the changes on the current connection."""
        if self.DEBUG==True: print '-- COMMIT --'
        self.connection.commit()

    def rollback(self):
        """rollback the changes on the current connection, aborting the transaction."""
        self.connection.rollback()

    def select(self, sql, vals=None, record_class=Record, recordset_class=RecordSet, cursor=None):
        """select from db and return the full result set.
        Required Arguments:
            sql: the SQL query as a string
        Optional/Named Arguments
            vals: any bound variables
            record_class: the class (itself) that the resulting records should be
        """
        c = cursor or self.cursor()
        self.execute(sql, vals=vals, cursor=c)

        # get a list of attribute names from the cursor.description
        attr_list = list()
        for r in c.description:
            attr_list.append(r[0])

        rs = recordset_class()                  # Populate a RecordSet (list) with the all resulting
        results = c.fetchall()
        for res in results:
            rec = record_class(self)            # whatever the record class is, include another instance
            for i in range(len(attr_list)):     # make each attribute dict-able by name
                rec[attr_list[i]] = res[i]
                if type(rec[attr_list[i]]) == str:  # always decode utf-8 into unicode
                    rec[attr_list[i]] = rec[attr_list[i]].decode('utf-8')
            rs.append(rec)                      # append the instance to the RecordSet

        if cursor is None:
            c.close()   # closing the cursor without committing rolls back the transaction.
        return rs

    def select_one(self, sql, vals=None, record_class=Record, cursor=None):
        """select one record from db
        Required Arguments:
            sql: the SQL query as a string
        Optional/Named Arguments:
            vals: any bound variables
            record_class: the class (itself) that the resulting records should be
        """
        c = cursor or self.cursor()
        self.execute(sql, vals, cursor=c)

        # get a list of attribute names from the cursor.description
        attr_list = list()
        for r in c.description:
            attr_list.append(r[0])
        res = c.fetchone()
        if res is None:
            rec = None
        else:
            rec = record_class(self)
            for i in range(len(attr_list)):
                rec[attr_list[i]] = res[i]
        if cursor is None:
            c.close()
        return rec

    def quote(self, attr):
        """returns the given attribute in a form that is insertable in the insert() and update() methods."""
        t = type(attr)
        if t == type(None): return u'NULL'
        elif t == datetime.datetime:    # datetime -- put it in quotes
            return u"'%s'" % unicode(attr)
        elif type(attr) in (unicode, str):
            return self._quote_str(attr)
        else:                           # boolean or number -- no quoting needed
            return unicode(attr).lower()

    def _quote_str(self, attr):
        """quote the attr string in a manner fitting the Database server, if known."""
        sn = self.servername().lower()
        if 'sqlserver' in sn or 'sqlite' in sn:
            # quote for sqlserver and sqlite: double '' to escape
            attr = "'" + re.sub("'", "''", attr) + "'"
        elif 'postgres' in sn:
            if type(attr) == unicode:
                attr = u"$$%s$$" % attr
            elif type(attr) == str:
                attr = u"$$%s$$" % unicode(attr, 'UTF-8')
        else:
            if type(attr) == unicode:
                attr = u"'%s'" % attr
            else:
                attr = u"'%s'" % unicode(attr, 'UTF-8')
        return attr

    def servername(self):
        """return a string that describes the database server being used"""
        if 'psycopg' in self.dba.__name__: return 'postgresql'
        elif 'sqlite' in self.dba.__name__: return 'sqlite'
        elif self.config and self.config.Database and self.config.Database.server:
            return self.config.Database.server
        elif 'adodbapi' in str(self.connection):
            return 'sqlserver'
        elif 'port=5432' in str(self.connection):
            return 'postgresql'
        elif 'port=3306' in str(self.connection):
            return 'mysql'
        else:
            return ''
        
    def table_exists(self, table_name):
        sn = self.servername().lower()
        if 'sqlite' in sn:
            return self.select_one("select * from sqlite_master where name=? and type='table' limit 1", (table_name,))
        elif 'mysql' in sn:
            return self.select_one("show tables like %s", (table_name,))
        else:   # postgresql and sqlserver both use the sql standard here.
            return self.select_one("select * from information_schema.tables where table_name=%s limit 1", (table_name,))


def doctests():
    """doctests for the amplitude.db module
>>> d = Database("dbname=test user=postgres password=pgsql")
>>> d.execute("create table table1 (name varchar primary key, email varchar not null unique);")
>>> d.execute("insert into table1 (name, email) values ('sah', 'sah@tyndale.com')")
>>> d.execute("insert into table1 (name, email) values ('sah', 'harrison@tbc.net')")
Traceback (most recent call last):
  ...
IntegrityError: duplicate key value violates unique constraint "table1_pkey"
<BLANKLINE>
>>> d.execute("drop table table1")
>>> d.connection.close()
    """

if __name__ == "__main__":
    import doctest
    doctest.testmod()
